"""This module is used to work with the CCHR segment files (initial data
files).
Exported Classes:
Exported Methods:
Exceptions:
ConvertError
"""
import cardsharp as cs
import os, re, json, shutil
from functools import wraps
from operator import itemgetter, attrgetter
from itertools import izip
import sys
from util import FORMAT_INFO, RUNDATETIME
from errors import *
from xlrd.biffh import XLRDError
from contextlib import closing
import MySQLdb
from configuration import segment_info, config
from util import *
__all__ = ['Segment', 'segment_list', 'segment_dict', 'all_segments']
segment_dict = dict()
segment_list = []
all_segments = []
for s in segment_info.itervalues():
segment_dict[s['label']] = {'label':s['label'], 'id' : s['id'], 'region_var' : s['region_var'], 'priority': s['priority']}
all_segments.append(s['label'])
for k, v in segment_dict.iteritems():
segment_list.append(v)
[docs]class Segment(object):
"""Segment object contains various methods to load, manipulate and save
CCHR segment data files.
Public Methods:
load_vars -- Loads the CCHR variable set as defined in the metadata.
get_index -- Takes a variable name and returns its index.
load_data -- Loads the segment data.
convert_vars -- Converts the types of all the variables to match the metadata file
get_data -- Returns the loaded dataset or None
add_new_vars -- Adds derived standardized varaibles to the dataset.
set_length_rules -- Sets length rules on all variables in the ds based on metadata file.
find_var_lengths -- Determines longest length for the original varaible in
the dataset and updates the metadata file with this information.
assign_var_lengths -- Assigns length rules to all the variables.
get_max_length -- Returns the highest varaible length in the dataset.
save -- Saves the dataset.
"""
def __init__(self, segment, phase, **opt):
self.verbose = opt['verbose']
if self.verbose:
print 'initializing Segment object'
self.name = segment
self.phase = phase
self.options = opt
self.length_rules = {}
self.precision_rules = {}
self.scale_rules = {}
self.fk_rules = {}
self.not_null_rules = {}
self.validate_rules = False
#add in suffix string for in_filter and out_filter
self.options['suffix'] = ''
for key in ['in_filter', 'out_filter']:
if key in opt and opt[key]:
if isinstance(opt[key], tuple):
out = '_'.join([opt[key][0],
str(opt[key][1])]).replace(' ', '')
else:
out = opt[key]
out = re.sub('\[|\]|,', '_', out)
self.options['suffix'] += ''.join(['_', key, '_', out])
#create placeholders
self.count_freqs = dict()
self.ds = None
if self.verbose:
print 'Segment object created'
#TODO: add try, except blocks on the variable information with robust
#error messages because this is a section that users interact with....
#or add the excel files to the web interface and add the validation
#to the web interface
[docs] def get_index(self, name, original = False):
"""Return the variable index of a supplied variable.
Arguments:
name -- The name of the variable who's index we want to find.
original -- Set to True to return the original index
"""
for row in self.var_info:
if row['name'] == name:
return row['index'] if not original else row['original']
def _drop_tables(self):
"""Helpr function to drop segment data tables in the database"""
opt = self.options
with closing(get_cnxn(opt['db_info']['name'],opt)) as cnx:
with closing(cnx.cursor()) as c:
c.execute('show tables')
response = c.fetchall()
for table in response:
if self.name == table[0]:
c.execute("DROP TABLE %s" % self.name)
elif self.name + 'x' == table[0]:
c.execute("DROP TABLE %s" % self.name + 'x')
[docs] def load_data(self):
"""Load the segment data.
If loading delimited data this is assumed to be the original segment
data so it is loading using the original variable names. If in_filter
was included as an option then will apply the filter variable function.
If limit was included in the options then will limit the rows loaded based
on the limit value.
"""
self.len_rules = {}
opt = self.options
#create the load options
state = '_' + opt.get('region').lower() if opt.get('region') else ''
options = {'source': os.path.join(opt['in_dir'], self.name + '%s%s' % (state,FORMAT_INFO[opt['in_format']])),
'format': opt['in_format'],
}
if opt.get('limit'):
options['limit'] = opt['limit']
if opt['in_format'] in ['del', 'text', 'csv']:
options['delimiter'] = opt['in_delimiter']
if opt['in_format'] == 'del':
options['escape_char'] = None
#if text and region specified or we are using processed data
#or converting staged data than has variable names, skip this row
if self.phase == 'convert' and not opt['original']:
options['skip'] = 1
options['var_names'] = self.orig_var_names if opt.get('original') else self.var_names
self.ds = cs.load(**options)
self.ds.wait()
#add count check for excel if > 6555? raise
[docs] def load_raw_data(self):
#TODO add user warning when rerunning stage that we will drop existing segment data when rerunning
#TODO use CONCURRENT to enable speedup
opt = self.options
self._drop_tables()
tbl_name = ''.join([self.name, 'x' if self.name != 'subject' else ''])
#stage the subject data into the subject table
if self.name == 'subject':
print "reading subject file and writing to Subject table...."
#load the relevant subject data
self.ds = cs.load(source=os.path.join(opt['in_dir'], 'subject.sav'), format='spss',
keep=[var.name if var.name != 'id' else 'casenum' for var in self.orig_variables])
#convert subject variables and set variable rules
self.add_new_vars()
self.convert_vars()
cs.wait()
self.ds.variables['casenum'].rename('id')
self.get_rules()
cs.wait()
self.set_rules()
cs.wait()
options = {'source': opt['db_info']['name'],
'dataset': self.name,
'format': 'mysql',
'user': opt['db_info']['user'],
'pwd': opt['db_info']['pass']
}
if self.verbose:
print 'start save...'
self.ds.save(**options)
cs.wait()
if self.verbose:
print 'done.'
print 'data saved'
print "complete."
with closing(get_cnxn(opt['db_info']['name'],opt)) as cnx:
with closing(cnx.cursor()) as c:
c.execute("SHOW TABLES FROM %s" % opt['db_info']['name'])
tables = [tbl[0] for tbl in c.fetchall()]
c.execute("""SELECT MAX(LENGTH(id)) FROM subject""")
casenum_len = [c.fetchone()[0]]
if self.name != 'subject':
#drop the table if it already exists
for tbl in [self.name, '%sx' % self.name]:
if tbl in tables:
c.execute("DROP TABLE %s" % tbl)
#create the table
vars = ','.join(['%s TEXT' % n for n in self.orig_var_names])
c.execute('''CREATE TABLE %s (id int(10) AUTO_INCREMENT, %s, PRIMARY KEY(id))''' % (tbl_name, vars))
#load the data
path = os.path.join(opt['in_dir'], '%s.%s' % (self.name, opt['in_format'])).replace('\\', '/')
num_vars = len(self.orig_var_names)
c.execute('''
LOAD DATA LOCAL INFILE '%s'
INTO TABLE %s
FIELDS TERMINATED BY %s ESCAPED BY ''
LINES TERMINATED BY '\r\n'
%s
SET %s
''' % (path,
tbl_name,
"'%s'" % opt.get('delimiter', '|'),
'(%s)' % ','.join(['@v%s' % x for x in xrange(num_vars)]),
'id=\N, %s' % ','.join(['%s = @v%s' % (name,i)
for i, name in enumerate(self.orig_var_names)])
)
)
cnx.commit()
#get casenum length
if self.name == 'arrest':
print "reading arrest file and writing to Arrestx table...."
#TODO after updating mysql driver make into tinyint
#TODO add column based on the disjoint of self.variables and self.orig_variables
new_vars = ['astate', 'asource', 'arrname','afed', 'cycnum', 'aseqnum', 'aori']
casenum_len.extend([self.var_dict[v][2] for v in new_vars])
c.execute('''
ALTER TABLE arrestx
ADD COLUMN casenum int(%i)
AFTER id,
ADD CONSTRAINT arrestx_casenum
FOREIGN KEY (casenum)
REFERENCES Subject(id),
ADD COLUMN astate int(%i)
AFTER arrstatex,
ADD CONSTRAINT arrestx_state
FOREIGN KEY (astate) REFERENCES region(id),
ADD COLUMN asource int(%i) AFTER rapstatex,
ADD CONSTRAINT arrestx_source
FOREIGN KEY (asource) REFERENCES source(id),
ADD COLUMN arrname nvarchar(%i) AFTER arrdatex,
ADD COLUMN afed int(%i) AFTER ynfederalx,
ADD CONSTRAINT arrestx_fed
FOREIGN KEY (afed) REFERENCES fed(id),
ADD COLUMN cycnum int(%i) AFTER cycnumx,
ADD COLUMN aseqnum int(%i) AFTER aseqnumx,
ADD COLUMN aori_invalid tinyint(1) AFTER arrnamex,
ADD COLUMN aori nvarchar(%i) AFTER aorix
''' % tuple(casenum_len)
)
elif self.name == 'sentence':
print "reading sentence file and writing to sentencex table...."
new_vars = ['cstate', 'csource', 'cfed', 'cycnum', 'cseqnum', 'cori']
casenum_len.extend([self.var_dict[v][2] for v in new_vars])
#TODO after updating mysql driver make into tinyint
c.execute('''ALTER TABLE sentencex
ADD COLUMN casenum int(%i)
AFTER id,
ADD CONSTRAINT sentencex_casenum
FOREIGN KEY (casenum)
REFERENCES Subject(id),
ADD COLUMN cstate int(%i)
AFTER crtstatex,
ADD CONSTRAINT sentencex_state
FOREIGN KEY (cstate)
REFERENCES region(id),
ADD COLUMN csource int(%i)
AFTER rapstatex,
ADD CONSTRAINT sentencex_source
FOREIGN KEY (csource)
REFERENCES source(id),
ADD COLUMN cfed int(%i)
AFTER ynfederalx,
ADD CONSTRAINT sentencex_fed
FOREIGN KEY (cfed)
REFERENCES fed(id),
ADD COLUMN cycnum int(%i)
AFTER cycnumx,
ADD COLUMN cseqnum int(%i)
AFTER cseqnumx,
ADD COLUMN cori nvarchar(%i)
AFTER crtnamex,
ADD COLUMN cori_invalid tinyint(1)
AFTER corix
''' % tuple(casenum_len)
)
elif self.name == 'supervision':
print "reading supervision file and writing to supervisionx table...."
casenum_len.extend([self.var_dict[v][2] for v in ['cycnum']])
c.execute('''ALTER TABLE supervisionx
ADD COLUMN casenum int(%i) AFTER id,
ADD CONSTRAINT supervisionx_casenum FOREIGN KEY (casenum) REFERENCES Subject(id),
ADD COLUMN cycnum int(%i) AFTER cycnumx
''' % tuple(casenum_len)
)
elif self.name == 'demographic':
print "reading demographic file and writing to demographicx table...."
#TODO after updating mysql driver make into tinyint
c.execute('''ALTER TABLE demographicx
ADD COLUMN casenum int(%i) AFTER id,
ADD CONSTRAINT demographicx_casenum FOREIGN KEY (casenum) REFERENCES Subject(id)
''' % tuple(casenum_len))
print "complete."
cnx.commit()
[docs] def load_staged_data(self):
if self.verbose:
print 'begin data load...'
opt = self.options
#TODO remove in_format from run command
options = {'source': opt['db_info']['name'],
'dataset': '%sx' % self.name,
'format': 'mysql',
'user': opt['db_info']['user'],
'pwd': opt['db_info']['pass'],
'load_as_null': ['',],
}
if opt.get('limit'):
options['sql_limit'] = '0, %s' % opt['limit']
#TODO change this to in_filter
if opt.get('in_where'):
if self.phase in ('pre_process', 'process') and re.match('state', opt.get('in_where')):
opt['in_where'] = re.sub('(?<!rap)state', '%sstate' % 'a' if (self.name == 'arrest') else 'cstate', opt['in_where'])
options['where'] = opt['in_where']
self.ds = cs.load(**options)
#TODO test to see if we need this wait
#cs.wait()
if self.verbose:
print '...start setting rules...'
if self.phase in ('process', 'convert'):
self.add_new_vars()
self.get_rules()
self.set_rules()
self.convert_vars()
cs.wait()
if self.verbose:
print '...finish setting rules...'
if self.verbose:
print 'data load complete.'
[docs] def load_processed_data(self, as_cache=False):
opt = self.options
#set the load options
options = {'source': opt['db_info']['name'],
'dataset': self.name,
'format': 'mysql',
'user': opt['db_info']['user'],
'pwd': opt['db_info']['pass'],
'load_as_null': ['',],
}
if as_cache:
options['select'] = self.orig_var_names
options['as_cache'] = True
options['cache_key'] = ['casenum']
#limit the number of rows loaded
if opt.get('limit'):
options['limit'] = opt['limit']
#specifiy the where clause
if opt.get('in_where'):
if self.phase == 'process':
opt['in_where'] = opt['in_where'].replace('state', 'astate' if self.name == 'arrest' else 'cstate')
options['where'] = opt['in_where']
#load the data into the memcache server
self.ds = cs.load(**options)
[docs] def create_recid_table(self):
#set the save options
options = {'source': self.options['recid_db_info']['name'],
'dataset': self.name,
'format': 'mysql',
'user': self.options['recid_db_info']['user'],
'pwd': self.options['recid_db_info']['pass'],
}
self.ds = cs.Dataset(self.variables)
self.get_rules()
self.set_rules()
#only save table if it does not already exist
self.ds.save(**options)
[docs] def add_new_vars(self):
"""Add new variables"""
if self.verbose:
print 'adding new vars...'
for row in self.var_info:
if self.name == 'subject' and row['original'] == False:
self.ds.variables.insert(int(row['index']), (row['name'], row['format']))
continue
#TODO make variables not dataset specific i.e. aori_invalid, cori_invalid = ori_invalid
if row['original'] is None and row['name'] not in ('id', 'casenum', 'cycnum', 'aori', 'cori', 'aori_invalid', 'cori_invalid',
'arrname', 'asource', 'csource', 'astate', 'cstate', 'afed', 'cfed',
'aseqnum', 'cseqnum'):
self.ds.variables.insert(int(row['index']), (row['name'], row['format']))
if self.verbose:
print 'done.'
[docs] def get_rules(self):
"""Get the rules."""
if self.verbose:
print 'start get_rules..'
for row in self.var_info:
if row.get('length') is not None:
self.length_rules[row['name']] = row['length']
if row.get('precision') is not None:
self.precision_rules[row['name']] = row['length'] + row['precision']
self.scale_rules[row['name']] = row['precision']
if row.get('fk') is not None:
self.fk_rules[row['name']] = row['fk']
if self.verbose:
print 'done.'
[docs] def set_rules(self):
"""Assigns rules to the variables associated with cardsharp dataset stored on self.ds.
This is not used during process command because we assign the rules during the data_load."""
if self.verbose:
print 'start set_rules...'
for rule_type, rule_name in [(self.length_rules, 'length'),
(self.precision_rules, 'precision'),
(self.scale_rules, 'scale'),
(self.fk_rules, 'fk')]:
for name, value in rule_type.iteritems():
if name in self.ds.variables:
if rule_name == 'length':
self.ds.variables[name].rules.length = value
elif rule_name == 'precision':
self.ds.variables[name].rules.precision = value
elif rule_name == 'scale':
self.ds.variables[name].rules.scale = value
elif rule_name == 'fk':
self.ds.variables[name].flags.foreign_key = ('%s_%s' % (self.name, value), ''.join([value, '(id)']))
if 'id' in self.ds.variables:
self.ds.variables['id'].flags.primary_key = 'ID'
self.ds.variables['id'].flags.auto_inc = False
if self.verbose:
print 'done.'
[docs] def get_max_length(self):
"""Return the highest variable length in the dataset."""
max = 0
for row in self.var_info:
max = row['length'] if row['length'] > max else max
return max
[docs] def convert_vars(self):
#convert if var_info format does not match current dataset format
for row in self.var_info:
if self.options.get('original') and row['original'] is None and self.phase == 'convert': #handle conversion for original data
continue
#remove x to convert from spss original file to del file
if self.options.get('original') and self.phase == 'convert' and row['name'].endswith('x'):
if row['format'] != self.ds.variables[row['name'][:-1]].format:
#TODO add catch for KeyError from cardsharp
self.ds.variables[row['name'][:-1]].convert(row['format'])
else:
if row['format'] != self.ds.variables[row['name']].format:
self.ds.variables[row['name']].convert(row['format'])
[docs] def drop_vars(self, format):
if format == 'mysql':
pass
[docs] def save(self):
"""Save the dataset.
Supported Keywords:
"""
opt = self.options
options = {'format': opt['out_format'], 'overwrite': True}
suffix = '' if not opt.get('suffix') else opt['suffix']
if opt['out_format'] in ['spss', 'mysql']:
options['max_string_width'] = self.get_max_length()
if opt['out_format'] in ['text', 'del', 'spss', 'excel']:
#add variable labels in spss files
if opt['out_format'] == 'spss':
for var, meta in self.metadata.iteritems():
self.ds.variables[var].label = meta['label']
self.ds.variables[var].value_labels = meta['value_labels']
#get the file extension
ext = FORMAT_INFO[opt['out_format']]
#get/create the output directory path
region = None
if opt.get('in_where'):
if len(re.findall('state[ ]*=[ ]*\d+', opt['in_where'])) == 1:
folder = get_state_key(re.search('(?<=state=)\d+', opt['in_where']).group())
region = folder
else:
folder = opt['in_where'].replace('=', '_')
else:
folder = 'all'
dir_path = os.path.normpath(os.path.join(opt['out_dir'], folder))
if not os.path.exists(dir_path):
os.makedirs(dir_path)
#set the save path
options['source'] = os.path.join(dir_path, ''.join([self.name, suffix,ext]))
#set the delimiter and turn off escaping if saving raw text
if opt['out_format'] in ['text', 'del']:
options['delimiter'] = opt['out_delimiter']
options['no_escape'] = True
elif opt['out_format'] in ['mysql',]:
options['overwrite'] = False
options['source'] = opt['db_info']['name']
options['user'] = opt['db_info']['user']
options['pwd'] = opt['db_info']['pass']
options['replace'] = True
if self.phase == 'pre_process':
options['dataset'] = '%sx' % self.name
else:
options['dataset'] = self.name
if self.name.lower() == 'arrest':
for v in ['achglit', 'achgdet', 'achgncic', 'aoffcode', ]:
self.ds.variables.rename(v + 'x', v)
options['drop'] = ['casenumx', 'arrstatex', 'rapstatex', 'ynfederalx', 'cycnumx',
'aseqnumx', 'arrdatex', 'arrnamex', 'aorix', 'offdatex', 'arrtypex',
'astatnumx', 'achgsvrx', 'atrknumx', 'aincx', 'arrcrfx', 'arrcefx',
'achgcntx', 'aactlitx', 'adispdatex', 'adisptypex', 'eventdate1x']
elif self.name.lower() == 'sentence':
for v in ['cchglit', 'cchgdet', 'cchgncic', 'coffcode']:
self.ds.variables.rename(v + 'x', v)
options['drop'] = ['casenumx', 'crtstatex', 'rapstatex', 'ynfederalx', 'cycnumx',
'cseqnumx', 'crtcmtx', 'cstatnumx', 'cchgsvrx', 'ctrknumx', 'cincx', 'crtcrfx',
'crtcefx', 'cchgcntx', 'cactlitx', 'cdispdatex', 'cdisptypex', 'fdispdatex',
'snttextx', 'cnflngx', 'cnflngminx', 'cnflngmaxx', 'susplngx', 'suspminx',
'suspmaxx', 'problingx', 'probminx', 'probmaxx', 'courtcostx', 'finex',
'restitutionx', 'crtnamex', 'corix', 'eventdate1x']
elif self.name.lower() == 'demographic':
options['drop'] = ['casenumx', 'demstatex', 'rapstatex', 'sexx', 'racex',
'pobx', 'yndeadx', 'dodx', 'hispanicx', 'ctzx', 'multistatex']
self.ds.save(**options)
cs.wait()
#copy the saved file to review after save
if opt['out_format'] == 'spss' and self.name in ('arrest', 'sentence', 'demographic'):
_region = region if region else 'all'
self.review_path = copy_to_review(self.name, _region, os.path.join(opt['out_dir'], _region, '%s.sav' % self.name))